import os
import re

def rename_txt_files_by_order(txt_dir, start_number=1):
    """
    将指定目录下的txt文件按照原始文件名从小到大的顺序重新命名
    
    Args:
        txt_dir (str): txt文件所在的目录路径
        start_number (int): 起始编号，默认为1
    """
    print("=" * 60)
    print("TXT文件重命名工具")
    print("=" * 60)
    print(f"目标目录: {txt_dir}")
    print(f"起始编号: {start_number}")
    print("-" * 60)
    
    # 检查目录是否存在
    if not os.path.exists(txt_dir):
        print(f"❌ 错误: 目录不存在 {txt_dir}")
        return
    
    # 获取所有txt文件
    txt_files = [f for f in os.listdir(txt_dir) if f.endswith('.txt')]
    
    if not txt_files:
        print(f"❌ 在 {txt_dir} 中没有找到txt文件")
        return
    
    print(f"📂 找到 {len(txt_files)} 个txt文件")
    print("-" * 60)
    
    # 按照文件名进行排序（自然排序，处理数字排序）
    def natural_sort_key(filename):
        """自然排序的key函数，正确处理文件名中的数字"""
        return [int(text) if text.isdigit() else text.lower() for text in re.split(r'(\d+)', filename)]
    
    txt_files.sort(key=natural_sort_key)
    
    # 显示排序后的文件列表
    print("排序后的文件列表:")
    for i, filename in enumerate(txt_files[:10]):  # 只显示前10个
        print(f"  {i+1}. {filename}")
    if len(txt_files) > 10:
        print(f"  ... 还有 {len(txt_files) - 10} 个文件")
    print("-" * 60)
    
    # 创建临时文件名映射，避免重命名冲突
    temp_mapping = {}
    renamed_count = 0
    
    # 第一步：将所有文件重命名为临时名称
    for idx, txt_file in enumerate(txt_files):
        old_file_path = os.path.join(txt_dir, txt_file)
        temp_name = f"temp_{idx}_{txt_file}"
        temp_file_path = os.path.join(txt_dir, temp_name)
        
        try:
            os.rename(old_file_path, temp_file_path)
            temp_mapping[temp_name] = start_number + idx
        except Exception as e:
            print(f"❌ 临时重命名失败 {txt_file}: {str(e)}")
            continue
    
    # 第二步：将临时文件重命名为最终名称
    for temp_name, final_number in temp_mapping.items():
        temp_file_path = os.path.join(txt_dir, temp_name)
        
        # 获取原始文件名
        original_name = temp_name.replace(f"temp_{list(temp_mapping.keys()).index(temp_name)}_", "")
        
        # 创建新的文件名
        new_file_name = f"{final_number}.txt"
        new_file_path = os.path.join(txt_dir, new_file_name)
        
        try:
            os.rename(temp_file_path, new_file_path)
            print(f"✅ 已将 {original_name} 重命名为 {new_file_name}")
            renamed_count += 1
        except Exception as e:
            print(f"❌ 最终重命名失败 {temp_name}: {str(e)}")
            # 如果失败，尝试恢复原名
            try:
                os.rename(temp_file_path, os.path.join(txt_dir, original_name))
            except:
                pass
    
    print("-" * 60)
    print(f"✅ 重命名完成! 成功重命名 {renamed_count}/{len(txt_files)} 个文件")
    print("=" * 60)

def main():
    """主函数 - 可以在这里配置参数"""
    
    # ===== 配置区域 - 修改这里的参数 =====
    CONFIG = {
        'txt_dir': r'F:\RDK_Tool_Train_Dataset\labels\train',  # txt文件目录路径
        'start_number': 1  # 起始编号
    }
    # =======================================
    
    # 也可以通过用户输入来设置参数
    print("TXT文件重命名工具")
    print("当前配置:")
    print(f"  TXT目录: {CONFIG['txt_dir']}")
    print(f"  起始编号: {CONFIG['start_number']}")
    
    user_input = input("\n是否使用当前配置? (y/n，直接回车默认为y): ").strip().lower()
    
    if user_input == 'n':
        new_dir = input("请输入TXT文件目录路径: ").strip()
        if new_dir:
            CONFIG['txt_dir'] = new_dir
        
        try:
            new_start = input(f"请输入起始编号 (默认{CONFIG['start_number']}): ").strip()
            if new_start:
                CONFIG['start_number'] = int(new_start)
        except ValueError:
            print("输入的起始编号无效，使用默认值")
    
    # 执行重命名
    rename_txt_files_by_order(CONFIG['txt_dir'], CONFIG['start_number'])

if __name__ == "__main__":
    main()